/*****************************************************
Copyright (C) 2020. All rights reserved.
File name     ：RedBlackTree.c
Version       ：v1.0
Author        ：zhengqijun
Date          ：2020-05-07
Function List ：
Description   ：Red Black Tree(Top-down)
******************************************************/

#include "RedBlackTree.h"
#include <stdio.h>   // for printf
#include <stdlib.h>  // for malloc
#include "FatalError.h"
#include "LinkedStack.h"
#include "LinkedQueue.h"

#define Infinity    (32767)   // 2^16-1
#define NegInfinity (-32768)  // -2^16

typedef enum {
    Red,
    Black,
} ColorType;

struct RedBlackNode {
    ElementType  Element;
    RedBlackTree Left;
    RedBlackTree Right;
    ColorType    Color;
};

// Static Function Lists
static Position Rotate(ElementType Item, Position Parent);
static void HandleReorient(ElementType Item, RedBlackTree T);
static void HandleReorientForRemove(ElementType Item, RedBlackTree T);
static Position SingleRotateWithLeft(Position K2);
static Position SingleRotateWithRight(Position K1);
static RedBlackTree MakeEmptyRec(RedBlackTree T);
static Position DoFind(ElementType Item, RedBlackTree T);
static Position DoFindMin(RedBlackTree T);
static Position DoFindMax(RedBlackTree T);
static void DoPreorderTraversal(RedBlackTree T);
static void DoInorderTraversal(RedBlackTree T);
static void DoPostorderTraversal(RedBlackTree T);
static void DoLevelorderTraversal(RedBlackTree T);
static void PrintDepth(ElementType Item, ColorType C, int depth);
static void DoPrintTree(RedBlackTree T, int depth);

// Static Variable Lists
static Position NullNode = NULL;  /* Needs initialization */

/* Initialization procedure */
RedBlackTree Initialize(void) {
    RedBlackTree T;

    if (NullNode == NULL) {
        NullNode = (Position)malloc(sizeof(struct RedBlackNode));
        if (NullNode == NULL) {
            FatalError("Out of space!!!");
        }

        NullNode->Left = NullNode;
        NullNode->Right = NullNode;
        NullNode->Color = Black;
        NullNode->Element = Infinity;
    }

    /* Create the header node */
    T = (RedBlackTree)malloc(sizeof(struct RedBlackNode));
    if (T == NULL) {
        FatalError("Out of space!!!");
    }

    T->Element = NegInfinity;
    T->Left = NullNode;
    T->Right = NullNode;
    T->Color = Black;

    return T;
}

static RedBlackTree MakeEmptyRec(RedBlackTree T) {
    if (T != NullNode) {
        MakeEmptyRec(T->Left);
        MakeEmptyRec(T->Right);
        free(T);
    }

    return NullNode;
}

RedBlackTree MakeEmpty(RedBlackTree T) {
    T->Right = MakeEmptyRec(T->Right);
    return T;
}

void Destroy(RedBlackTree T) {
    MakeEmpty(T);
    free(T);
}

int IsEmpty(RedBlackTree T) {
    return (T == NULL) || (T == NullNode) || (T->Right == NullNode);
}

static Position DoFind(ElementType Item, RedBlackTree T) {
    while (T != NullNode) {
        if (Item < T->Element) {
            T = T->Left;
        } else if (Item > T->Element) {
            T = T->Right;
        } else {
            break;
        }
    }

    return T;
}

Position Find(ElementType Item, RedBlackTree T) {
    return DoFind(Item, T->Right);
}

static Position DoFindMin(RedBlackTree T) {
    if (T != NullNode) {
        while(T->Left != NullNode) {
            T = T->Left;
        }
    }

    return T;
}

Position FindMin(RedBlackTree T) {
    return DoFindMin(T->Right);
}

static Position DoFindMax(RedBlackTree T) {
    if (T != NullNode) {
        while (T->Right != NullNode) {
            T = T->Right;
        }
    }

    return T;
}

Position FindMax(RedBlackTree T) {
    return DoFindMax(T->Right);
}

/* This function can be called only if K2 has a left child */
/* Perform a rotate between a node (K2) and its left child */
/* Update heights, then return new root */
static Position SingleRotateWithLeft(Position K2) {
    Position K1;

    K1 = K2->Left;
    K2->Left = K1->Right;
    K1->Right = K2;

    return K1;  /* New root */
}

/* This function can be called only if K1 has a right child */
/* Perform a rotate between a node (K1) and its right child */
/* Update heights, then return new root */
static Position SingleRotateWithRight(Position K1) {
    Position K2;

    K2 = K1->Right;
    K1->Right = K2->Left;
    K2->Left = K1;

    return K2;  /* New root */
}

/* Perform a rotation at node X */
/* (whose parent is passed as a parameter) */
/* The child is deduced by examining Item */
static Position Rotate(ElementType Item, Position Parent) {
    if (Item < Parent->Element) {
        return Parent->Left = Item < Parent->Left->Element ?
            SingleRotateWithLeft(Parent->Left) :
            SingleRotateWithRight(Parent->Left);
    } else {
        return Parent->Right = Item < Parent->Right->Element ?
            SingleRotateWithLeft(Parent->Right) :
            SingleRotateWithRight(Parent->Right);
    }
}

static Position X, P, GP, GGP;

static void HandleReorient(ElementType Item, RedBlackTree T) {
    X->Color = Red;        /* Do the color flip */
    X->Left->Color = Black;
    X->Right->Color = Black;

    /* Have to rotate */
    if (P->Color == Red) {
        GP->Color = Red;
        if ((Item < GP->Element) != (Item < P->Element)) {
            P = Rotate(Item, GP);  /* Start double rotate */
        }

        X = Rotate(Item, GGP);
        X->Color = Black;
    }

    T->Right->Color = Black;  /* Make root black */
}

RedBlackTree Insert(ElementType Item, RedBlackTree T) {
    X = T;
    P = T;
    GP = T;
    NullNode->Element = Item;

    /* Descend down the tree */
    while (X->Element != Item) {
        GGP = GP;
        GP = P;
        P = X;
        if (Item < X->Element) {
            X = X->Left;
        } else {
            X = X->Right;
        }

        if (X->Left->Color == Red && X->Right->Color == Red) {
            HandleReorient(Item, T);
        }
    }

    if (X != NullNode) {
        return NullNode;  /* Duplicate */
    }

    X = (Position)malloc(sizeof(struct RedBlackNode));
    if (X == NULL) {
        FatalError("Out of space!!!");
    }

    X->Element = Item;
    X->Left = NullNode;
    X->Right = NullNode;

    /* Attach to its parent */
    if (Item < P->Element) {
        P->Left = X;
    } else {
        P->Right = X;
    }

    HandleReorient(Item, T);  /* Color it red; maybe rotate */

    return T;
}

static void HandleReorientForRemove(ElementType Item, RedBlackTree T) {
    Position B;
    Position R;
    Position Tmp;

    if (P->Left == X) {
        B = P->Right;
    } else {
        B = P->Left;
    }

    if (X->Left->Color == Black && X->Right->Color == Black) {
        if (B->Left->Color == Black && B->Right->Color == Black) {
            /* Do the color flip */
            X->Color = Red;
            P->Color = Black;
            B->Color = Red;
        } else if (B->Left->Color == Red) {
            R = B->Left;
            P->Color = Black;
            X->Color = Red;

            if (GP->Left == P) {
                if (P->Left == B) {
                    B->Color = Red;
                    R->Color = Black;
                    GP->Left = SingleRotateWithLeft(P);

                    GGP = GP;
                    GP = B;
                } else {
                    P->Right = SingleRotateWithLeft(B);
                    GP->Left = SingleRotateWithRight(P);

                    GGP = GP;
                    GP = R;
                }
            } else {
                if (P->Left == B) {
                    B->Color = Red;
                    R->Color = Black;
                    GP->Right = SingleRotateWithLeft(P);

                    GGP = GP;
                    GP = B;
                } else {
                    P->Right = SingleRotateWithLeft(B);
                    GP->Right = SingleRotateWithRight(P);

                    GGP = GP;
                    GP = R;
                }
            }
        } else {
            R = B->Right;
            P->Color = Black;
            X->Color = Red;

            if (GP->Left == P) {
                if (P->Left == B) {
                    P->Left = SingleRotateWithRight(B);
                    GP->Left = SingleRotateWithLeft(P);

                    GGP = GP;
                    GP = R;
                } else {
                    B->Color = Red;
                    R->Color = Black;
                    GP->Left = SingleRotateWithRight(P);

                    GGP = GP;
                    GP = B;
                }
            } else {
                if (P->Left == B) {
                    P->Left = SingleRotateWithRight(B);
                    GP->Right = SingleRotateWithLeft(P);

                    GGP = GP;
                    GP = R;
                } else {
                    B->Color = Red;
                    R->Color = Black;
                    GP->Right = SingleRotateWithRight(P);

                    GGP = GP;
                    GP = B;
                }
            }
        }
    } else if (X->Left->Color == Red) {
        if (Item < X->Element) {
            Tmp = X->Left;
        } else if (Item > X->Element) {
            Tmp = X->Right;
        } else {
            return;
        }

        if (Tmp->Color == Red) {
            return;
        } else {
            X->Color = Red;
            Tmp->Color = Black;
            if (P->Left == X) {
                P->Left = SingleRotateWithLeft(X);
            } else {
                P->Right = SingleRotateWithLeft(X);
            }

            GGP = GP;
            GP = P;
            P = Tmp;
        }
    } else if (X->Right->Color == Red || X->Left->Color == Red) {
        if (Item < X->Element) {
            Tmp = X->Left;
        } else if (Item > X->Element) {
            Tmp = X->Right;
        } else {
            return;
        }

        if (Tmp->Color == Red) {
            return;
        } else {
            X->Color = Red;
            Tmp->Color = Black;
            if (P->Left == X) {
                P->Left = SingleRotateWithRight(X);
            } else {
                P->Right = SingleRotateWithRight(X);
            }

            GGP = GP;
            GP = P;
            P = Tmp;
        }
    }
}

RedBlackTree Remove(ElementType Item, RedBlackTree T) {
    Position Tmp;

    X = T;
    P = T;
    GP = T;
    NullNode->Element = Item;

    /* Make root red */
    T->Right->Color = Red;

    /* Descend down the tree */
    while (X != NullNode) {
        GGP = GP;
        GP = P;
        P = X;
        if (Item < X->Element) {
            X = X->Left;
        } else if (Item > X->Element) {
            X = X->Right;
        } else {

            if (X->Right != NullNode) {
                Tmp = DoFindMin(X->Right);

                X->Element = Tmp->Element;
                Tmp->Element = NullNode->Element;

                X = X->Right;
            } else if (X->Left != NullNode) {
                Tmp = DoFindMax(X->Left);

                X->Element = Tmp->Element;
                Tmp->Element = NullNode->Element;

                X = X->Left;
            } else {  /* Leaf node */
                if (X->Color == Red) {
                    if (GP->Left == X) {
                        GP->Left = NullNode;
                    } else {
                        GP->Right = NullNode;
                    }

                    Tmp = X;
                    X = NullNode;
                    free(Tmp);
                    break;
                } else {
                    GP = GGP;
                    P = GP;
                }
            }
        }

        if (P->Color == Red && P->Left->Color == Black && P->Right->Color == Black) {
            HandleReorientForRemove(Item, T);
        }
    }

    /* Make root black */
    T->Right->Color = Black;

    return T;
}

Position Left(Position P) {
    return P->Left;
}

Position Right(Position P) {
    return P->Right;
}

ElementType Retrieve(Position P) {
    return P->Element;
}

// Recursive implement
// static void DoPreorderTraversal(RedBlackTree T) {
//     if (T != NullNode) {
//         printf("%d ", T->Element);
//         DoPreorderTraversal(T->Left);
//         DoPreorderTraversal(T->Right);
//     }
// }

// Non-recursive implement
static void DoPreorderTraversal(RedBlackTree T) {
    if (T != NullNode) {
        LinkedStack S;
        Position P;

        S = CreateStack();

        P = T;
        while (P != NullNode || !IsEmptyStack(S)) {
            while (P != NullNode) {
                printf("%d ", P->Element);
                Push(P, S);
                P = P->Left;
            }

            P = TopAndPop(S);
            P = P->Right;
        }
    }
}

/* Print the tree, watch out for NullNode, */
/* and skip header */
void PreorderTraversal(RedBlackTree T) {
    DoPreorderTraversal(T->Right);
}

// Recursive implement
// static void DoInorderTraversal(RedBlackTree T) {
//     if (T != NullNode) {
//         DoInorderTraversal(T->Left);
//         printf("%d ", T->Element);
//         DoInorderTraversal(T->Right);
//     }
// }

// Non-recursive implement
static void DoInorderTraversal(RedBlackTree T) {
    if (T != NullNode) {
        LinkedStack S;
        Position P;

        S = CreateStack();

        P = T;
        while (P != NullNode || !IsEmptyStack(S)) {
            while (P != NullNode) {
                Push(P, S);
                P = P->Left;
            }

            P = TopAndPop(S);
            printf("%d ", P->Element);
            P = P->Right;
        }
    }
}

/* Print the tree, watch out for NullNode, */
/* and skip header */
void InorderTraversal(RedBlackTree T) {
    DoInorderTraversal(T->Right);
}

// Recursive implement
// static void DoPostorderTraversal(RedBlackTree T) {
//     if (T != NullNode) {
//         DoPostorderTraversal(T->Left);
//         DoPostorderTraversal(T->Right);
//         printf("%d ", T->Element);
//     }
// }

// Non-recursive implement
static void DoPostorderTraversal(RedBlackTree T) {
    if (T != NullNode) {
        LinkedStack S;
        Position Pre;
        Position Cur;

        S = CreateStack();

        Cur = T;
        Pre = NullNode;
        while (Cur != NullNode || !IsEmptyStack(S)) {
            if (Cur != NullNode) {
                Push(Cur, S);
                Cur = Cur->Left;
            } else {
                Cur = Top(S);
                if (Cur->Right != NullNode && Cur->Right != Pre) {
                    Cur = Cur->Right;
                } else {
                    Pop(S);
                    printf("%d ", Cur->Element);
                    Pre = Cur;
                    Cur = NullNode;
                }
            }
        }
    }
}

/* Print the tree, watch out for NullNode, */
/* and skip header */
void PostorderTraversal(RedBlackTree T) {
    DoPostorderTraversal(T->Right);
}

static void DoLevelorderTraversal(RedBlackTree T) {
    if (T != NullNode) {
        LinkedQueue Q;

        Q = CreateQueue();

        EnQueue(T, Q);

        while (!IsEmptyQueue(Q)) {
            if (Front(Q)->Left != NullNode) {
                EnQueue(Front(Q)->Left, Q);
            }

            if (Front(Q)->Right != NullNode) {
                EnQueue(Front(Q)->Right, Q);
            }

            printf("%d ", Retrieve(FrontAndDequeue(Q)));
        }
    }
}

/* Print the tree, watch out for NullNode, */
/* and skip header */
void LevelorderTraversal(RedBlackTree T) {
    DoLevelorderTraversal(T->Right);
}

static void PrintDepth(ElementType Item, ColorType C, int depth) {
    while (depth != 0) {
        printf("    ");
        --depth;
    }

    if (C == Red) {
        printf("\033[1;31m%d\033[0m\n", Item);
    } else if (C == Black) {
        printf("\033[1;36m%d\033[0m\n", Item);
    } else {
        Error("Error color");
    }
}

static void DoPrintTree(RedBlackTree T, int depth) {
    if (T == NullNode) {
        printf("RedBlackTree is empty!\n");
        return;
    }

    if (Right(T) != NullNode) {
        DoPrintTree(Right(T), depth + 1);
    }

    PrintDepth(Retrieve(T), T->Color, depth);

    if (Left(T) != NullNode) {
        DoPrintTree(Left(T), depth + 1);
    }
}

/* Print the tree, watch out for NullNode, */
/* and skip header */
void PrintTree(RedBlackTree T, int depth) {
    DoPrintTree(T->Right, depth);
}

// EOF

